from glob import glob

import pandas as pd
import polars as pl
import numpy as np
from tqdm import tqdm

voc = pd.read_pickle("total_voc_new.pkl")["voc"]
voc = ["<|sos|>", "<|user|>", "<|agent|>", "<|pad|>", "<|history|>", "<|unk|>", "<|end|>", "<|next|>"] + sorted(
    set(voc))
voc_size = len(voc) - 8
voc1 = pd.read_pickle("total_voc_new.pkl")["voc1"]
voc1 = pd.DataFrame({"voc": voc1.keys(), "count": voc1.values()})
voc1 = voc1[~voc1["voc"].isin(voc)]
voc1 = voc1.sort_values("count", ascending=False)
voc1["voc_id"] = np.array(list(range(len(voc1)))) % voc_size + 8
voc1["voc_num"] = np.array(list(range(len(voc1)))) // voc_size + 9

dl = pl.DataFrame(
    {"voc": voc + voc1["voc"].values.tolist(), "voc_id": list(range(len(voc))) + voc1["voc_id"].values.tolist(),
     "voc_num": [8] * len(voc) + voc1["voc_num"].values.tolist()})

paths = glob("D:/SamOutV2/SamOutV2/sft/*")
total_tokens = 0
for path in tqdm(paths):
    path_data = pd.read_pickle(path)
    for j in range(0,len(path_data),10000):
        data=path_data[j:j+10000]
        data_new = []
        for i in data:
            if len(i)>512:

                data_new.append(i[:512] + ["<|end|>"])
            else:
                data_new.append(i + ["<|end|>"]+(512-len(i))*["<|pad|>"])
        data = pd.DataFrame(data_new)
        data = data.fillna("<|pad|>")
        data = pd.DataFrame({"voc": data.values[:, :512].reshape(-1)})
        data = pl.from_pandas(data)
        data = data.join(dl, on="voc", how="left")

        data = data["voc_id", "voc_num"].to_numpy().reshape([-1, 512, 2]).tolist()
        data_list = []
        for one in tqdm(data):
            current_voc = 7
            new_token_list = []
            for i, (two, two_id) in enumerate(one):

                if two_id != current_voc:
                    current_voc = two_id
                    new_token_list += [7, current_voc]
                new_token_list.append(two)
            data_list.append(new_token_list[:512])


        pd.to_pickle(data_list, path.replace("sft", "sft_id").replace("data_single",str(j)), compression="zip")
